import io
from datetime import date, datetime
from typing import List, Union, Optional, Dict
import struct
import peewee
from emi.bar.data import Interval
from emi.bar.storage.model import BaseDBModel
from emi.core.DBSource import DBSource
from emi.trade.data import DailyResult, TradeData, Direction, Offset,OpType
from emi.util import NumUtils
from emi.io import write_float_compress, read_float_compress, write_int_compress, read_int_compress, write_str_list, \
    read_str_list

class DailyResultModel(BaseDBModel):

    day_token:int = peewee.SmallIntegerField(primary_key=True)
    time:date = peewee.DateField(index=True)
    balance:float = peewee.DoubleField(null=False)
    commission:float = peewee.DoubleField(null=False)
    balance_hold:float = peewee.DoubleField(null=False)

    trades_bytes: bytes = peewee.BlobField(null=True)
    extra_bytes:bytes = peewee.BlobField(null=True)

    save_type:int = peewee.IntegerField(index=True)


    def to_dict(self):
        return self.__data__

    @staticmethod
    def from_data(bar: DailyResult,type:int):
        model = DailyResultModel()
        model.day_token = bar.day_token
        model.time = bar.time.date()
        model.balance = bar.balance
        model.commission = bar.commission
        model.balance_hold = bar.balance_hold
        model.save_type = type
        model.trades_bytes = DailyResultModel._to_bytes(bar.tradeDatas)
        return model

    def to_bar(self) -> DailyResult:
        bar = DailyResult(
            day_token=self.day_token,
            time= datetime.fromordinal(self.time.toordinal()),
            balance=self.balance,
            commission=self.commission,
            balance_hold=self.balance_hold,
            tradeDatas= DailyResultModel._to_trades(self.trades_bytes)
        )
        return bar

    @staticmethod
    def _to_bytes(trades: List[TradeData]) -> bytes:
        if not trades:
            return None
        output = io.BytesIO()
        output.write(struct.pack("H", len(trades)))
        for trade in trades:
            str_list = [trade.symbol,
                        trade.orderid,
                        trade.direction.value,
                        trade.offset.value,
                        trade.op_type.value,
                        trade.msg
                        ]  ## 6
            float_list = [trade.price, trade.commission, trade.trade_money]
            int_list = [trade.day_token, int(trade.volume)]

            output.write(struct.pack("d", trade.time.timestamp()))
            write_int_compress(output, int_list)
            write_float_compress(output, float_list)
            write_str_list(output, str_list)
        return output.getvalue()

    @staticmethod
    def _to_trades(data: bytes) -> List[TradeData]:
        if not data:
            return []
        offset = 0
        trade_size = struct.unpack_from("H", data, offset)[0]
        offset += 2
        orders = []
        for i in range(0, trade_size):
            time = datetime.fromtimestamp(struct.unpack_from("d", data, offset)[0])
            offset += 8
            int_list, offset = read_int_compress(data, offset)
            float_list, offset = read_float_compress(data, offset)
            str_list, offset = read_str_list(data, offset)
            orders.append(TradeData(
                symbol=str_list[0], orderid=str_list[1], direction=Direction(str_list[2]),
                offset=Offset(str_list[3]), op_type=OpType(str_list[4]), msg=str_list[5],
                price=float_list[0], commission=float_list[1], trade_money=float_list[2],
                day_token=int_list[0], volume=int_list[1],
                time=time
            ))
        return orders

class DailyResultStorage:

    def __init__(self, source: DBSource, table_name=None):
        # indexes_desc = (
        #     (("dimen_value", "interval", "datetime"), True),
        # )
        if table_name is None:
            table_name = 'daily_result'
        self.source = source.init_peewee(table_name, DailyResultModel, indexes=None, ignore_inited=True)

    def save_bars(self, bars: List[DailyResult],type:int):
        """
        保存Bar数据。
        """
        ds = [DailyResultModel.from_data(i,type) for i in bars]
        with self.source.use_peewee() as db_model:
            db_model.save_all(ds)


    def get_newest_day(self)->datetime:
        with self.source.use_peewee() as db_model:
            s = db_model.select(db_model.time).order_by(db_model.datetime.desc()).first()
            if s:
                date = s.time
                return datetime(year=date.year, month=date.month, day=date.day)
            return None

    def get_day_list(self) -> List[date]:
        with self.source.use_peewee() as db_model:
            s = db_model.select(db_model.time).order_by(db_model.time.asc()).distinct()
            date_list = [db_bar.time for db_bar in s]
        return date_list

    def size(self) -> int:
        with self.source.use_peewee() as db_model:
            return db_model.select().count()

    def clear(self):
        with self.source.use_peewee() as db_model:
            db_model.drop_table()
            db_model.create_table()

    def get_newest(self)->DailyResult:
        today_dailys = self.get_all(type= 0)
        if today_dailys:
            return today_dailys[-1]
        return None

    def get_all(self,type:int) -> List[DailyResult]:
        """
        返回指定日期的量化数据。
        :param dimen:
        :param interval:
        :param tag:
        :return:
        """

        with self.source.use_peewee() as db_model:

            s = (
                db_model.select().where(db_model.save_type==type)
                    .order_by(db_model.time)
            )
            if s:
                return [db_bar.to_bar() for db_bar in s]
        return []
